from paddleocr import PaddleOCR, draw_ocr
from PIL import Image


def test_paddleocr():
    ocr = PaddleOCR(use_angle_cls=True, lang="ch")
    img_path = r'./data/test.jpg'
    result = ocr.ocr(img_path, cls=True)
    result = result[0]
    for line in result:
        print(line)

    # 显示结果
    image = Image.open(img_path).convert('RGB')
    boxes = [line[0] for line in result]
    txts = [line[1][0] for line in result]
    scores = [line[1][1] for line in result]
    im_show = draw_ocr(image, boxes, txts, scores)
    # im_show = draw_ocr(image, boxes)
    im_show = Image.fromarray(im_show)
    im_show.save(r'./data/test_res.jpg')


def test_tableocr():
    import os
    import cv2
    from paddleocr import PPStructure, draw_structure_result, save_structure_res

    table_engine = PPStructure(show_log=True)
    save_folder = r'./data'
    img_path = r'./data/tmp/image6.emf'
    font_path = r'./fonts/simfang.ttf'
    img = cv2.imread(img_path)
    result = table_engine(img)
    save_structure_res(result, save_folder, os.path.basename(img_path).split('.')[0])

    for line in result:
        line.pop('img')
        print(line)

    from PIL import Image
    image = Image.open(img_path).convert('RGB')
    im_show = draw_structure_result(image, result, font_path=font_path)
    im_show = Image.fromarray(im_show)
    im_show.save(r'./data/test_res.jpg')


if __name__ == '__main__':
    # test_paddleocr()
    test_tableocr()
